package com.pg.agent.engine.service.utils;

import cn.hutool.core.util.ObjectUtil;
import com.alibaba.fastjson.JSONObject;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.UUID;

@Component
@Slf4j
public class MilvusUtils {

  /*  @Resource
    private MilvusServiceClient milvusServiceClient;

    @Resource
    private Producer producer;

    @PostConstruct
    public void init() {
        milvusServiceClient.loadCollection(
                LoadCollectionParam.newBuilder()
                        .withCollectionName("chat_embedding")
                        .build());
    }

    public String getRelevantHistory(List<Double> embedding, String expr, int topK, String agentName, Integer maxHistory) {
        try {
            final String SEARCH_PARAM = "{\"nprobe\":" + maxHistory + "}";

            List<Float> floatList = new ArrayList<>();
            for (Double d : embedding) {
                float f = Float.parseFloat(Double.toString(d));
                floatList.add(f);
            }
            List<List<Float>> lists = new ArrayList<>();
            lists.add(floatList);

            List<String> id = Arrays.asList("send_name", "send_role", "content");
            SearchParam searchParam = SearchParam.newBuilder()
                    .withCollectionName("chat_embedding")
                    .withConsistencyLevel(ConsistencyLevelEnum.STRONG)
                    .withOutFields(id)
                    .withTopK(topK)
                    .withVectors(lists)
                    .withExpr(expr)
                    .withVectorFieldName("content_vector")
                    .withParams(SEARCH_PARAM)
                    .build();
            R<SearchResults> respSearch = milvusServiceClient.search(searchParam);
            if (respSearch.getStatus() == R.Status.Success.getCode()) {
                SearchResults resp = respSearch.getData();
                if (!resp.hasResults()) { //判断是否查到结果
                    return null;
                }
                SearchResultsWrapper wrapperSearch = new SearchResultsWrapper(resp.getResults());
                List<String> sendNames = (List<String>) wrapperSearch.getFieldData("send_name", 0);
                List<String> sendRoles = (List<String>) wrapperSearch.getFieldData("send_role", 0);
                List<String> content = (List<String>) wrapperSearch.getFieldData("content", 0);
                StringBuilder chatHistoryStr = new StringBuilder();
                for (int i = 0; i < sendNames.size(); i++) {
                    String name = sendNames.get(i);
                    if (ObjectUtil.equal(agentName, name)) {
                        name = "你";
                    }
                    chatHistoryStr.append(String.format(PromptConstant.SINGLE_MESSAGE, name, sendRoles.get(i), content.get(i)));
                }
                return String.format(PromptConstant.CHAT_TOTAL_MESSAGE, chatHistoryStr.toString());
            }
        } catch (Exception e) {
            log.info("查询数据失败：{}", e.getMessage());
        }
        return null;
    }

    public void insertData(Long userId, String sessionId, Long sopId, String stateName,
                           String sendName, String sendRole, String content, List<Double> embedding, Long saasAppId, Integer exeType, Integer type) {
        try {
            List<InsertParam.Field> fields = new ArrayList<>();
            fields.add(new InsertParam.Field("user_id", Collections.singletonList(userId)));
            fields.add(new InsertParam.Field("session_id", Collections.singletonList(sessionId)));
            fields.add(new InsertParam.Field("sop_id", Collections.singletonList(sopId)));
            fields.add(new InsertParam.Field("state_name", Collections.singletonList(stateName)));
            fields.add(new InsertParam.Field("send_name", Collections.singletonList(sendName)));
            fields.add(new InsertParam.Field("send_role", Collections.singletonList(sendRole)));
            fields.add(new InsertParam.Field("word_count", Collections.singletonList(Long.valueOf(content.length()))));
            fields.add(new InsertParam.Field("content", Collections.singletonList(content)));

            List<Float> floatList = new ArrayList<>();
            if (embedding != null) {
                for (Double d : embedding) {
                    float f = Float.parseFloat(Double.toString(d));
                    floatList.add(f);
                }
            }

            List<List<Float>> lists = new ArrayList<>();
            lists.add(floatList);
            fields.add(new InsertParam.Field("content_vector", lists));

            InsertParam insertParam = InsertParam.newBuilder()
                    .withCollectionName("chat_embedding")
                    .withFields(fields)
                    .build();
            //插入数据
            milvusServiceClient.insert(insertParam);
            JSONObject jsonObject = new JSONObject();
            jsonObject.put("userId", userId);
            jsonObject.put("sopId", sopId);
            jsonObject.put("sessionId", sessionId);
            jsonObject.put("stateName", stateName);
            jsonObject.put("sendName", sendName);
            jsonObject.put("sendRole", sendRole);
            jsonObject.put("wordCount", content.length());
            jsonObject.put("content", content);
            jsonObject.put("createTime", System.currentTimeMillis());
            jsonObject.put("saasAppId", saasAppId);
            jsonObject.put("exeType", exeType);
            jsonObject.put("type", type);
            producer.sendMsg("ENGINE_TEMPLATE", "chat_history", UUID.randomUUID().toString(), JSONObject.toJSONString(jsonObject));
        } catch (Exception e) {
            log.info("插入数据报错：{}", e.getMessage());
        }
    }

    public void insertMemory(ManagerObject managerObject, InPutMessage message, Memory memory) {
        {
            try {
                List<InsertParam.Field> fields = new ArrayList<>();
                fields.add(new InsertParam.Field("user_id", Collections.singletonList(managerObject.getUserId())));
                fields.add(new InsertParam.Field("session_id", Collections.singletonList(managerObject.getSessionId())));
                fields.add(new InsertParam.Field("sop_id", Collections.singletonList(message.getSop().getId())));
                fields.add(new InsertParam.Field("state_name", Collections.singletonList(message.getState() == null ? "初始":message.getState().getName())));
                fields.add(new InsertParam.Field("send_name", Collections.singletonList(memory.getSendName())));
                fields.add(new InsertParam.Field("send_role", Collections.singletonList(memory.getSendRole()== null ?memory.getSendName() : memory.getSendRole())));
                fields.add(new InsertParam.Field("word_count", Collections.singletonList(Long.valueOf(memory.getContent().length()))));
                fields.add(new InsertParam.Field("content", Collections.singletonList(memory.getContent())));

                List<Float> floatList = new ArrayList<>();
                if (memory.getChatEmbeddings() != null) {
                    for (Double d : memory.getChatEmbeddings()) {
                        float f = Float.parseFloat(Double.toString(d));
                        floatList.add(f);
                    }
                }

                List<List<Float>> lists = new ArrayList<>();
                lists.add(floatList);
                fields.add(new InsertParam.Field("content_vector", lists));

                InsertParam insertParam = InsertParam.newBuilder()
                        .withCollectionName("chat_embedding")
                        .withFields(fields)
                        .build();
                //插入数据
                milvusServiceClient.insert(insertParam);
//                JSONObject jsonObject = new JSONObject();
//                jsonObject.put("userId", userId);
//                jsonObject.put("sopId", sopId);
//                jsonObject.put("sessionId", sessionId);
//                jsonObject.put("stateName", stateName);
//                jsonObject.put("sendName", sendName);
//                jsonObject.put("sendRole", sendRole);
//                jsonObject.put("wordCount", content.length());
//                jsonObject.put("content", content);
//                jsonObject.put("createTime", System.currentTimeMillis());
//                jsonObject.put("saasAppId", saasAppId);
//                jsonObject.put("exeType", exeType);
//                jsonObject.put("type", type);
//                producer.sendMsg("ENGINE_TEMPLATE", "chat_history", UUID.randomUUID().toString(), JSONObject.toJSONString(jsonObject));
            } catch (Exception e) {
                e.printStackTrace();
                log.info("插入数据报错：{}", e.getMessage());
            }
        }
    }*/
}
